############# CES OLY ESTIMATION ##############

#multi-market
ces.oly.mm.update <- function(eta,pre,nu=0.3){
  mc <- pre$cit
  tau <- pre$cit/pre$mc
  s <- pre$s
  s.f <- pre$s.f
  i <- pre$i
  f <- pre$f
  m <- pre$m
  n <- pre$n
  epsilon <- eta+1 - eta*s.f
  markup <- epsilon/(epsilon-1)
  # estimate the demand and quality coefficient ("CES" and "B")
  ols.results <<- lm(log(s)+eta*log(markup) ~ log(tau) + factor(m) + factor(n)) 
  #
  cost.elas <- ols.results$coefficients[2]
  eta.star <- as.numeric(-cost.elas)
  return(nu*eta.star +(1-nu)*eta)
}


#Monte-Carlo of the estimation
monte.oly.mm.func <- function(repi,settings,v=1,eta.ces.oly.init) { # monte carlo stage
  data.exog.mm <- DGP.exog(settings=settings,v=v,grid=subset(tau.grid.allv,vtau==v,select=-vtau)) #gen exog data
  eqbm.pre.mm  <- as.data.table(DGP.endog(parallel=T,v=v,tar.change=0,data.exog = data.exog.mm))
  eqbm.pre.mm[, s.f := sum(s), by=c("f","n")]
  eqbm.pre.mm[, dyad := paste(i,"-",n) ]
  FPIresults <- fpiter(eta.ces.oly.init,fixptfn=ces.oly.mm.update,pre=eqbm.pre.mm,nu=1) #estimate eta.ces.oly
  se.eta.ces.oly <- sqrt(vcov(ols.results)[2,2]) # ols.results is created in ces.oly.mm.update
  se.eta.ces.oly.clust <- sqrt(vcovCR(ols.results, cluster = eqbm.pre.mm$dyad, type = "CR2")[2,2])
  eqbm.pre.mm.with.ces <- as.data.table(estapprox.mm(eqbm.pre.mm)) #add ces approximation results
  #output
  return(data.frame(repi=repi,v=v,mu.eta=exp(settings$mu.a[v]),pelas=mean(eqbm.pre.mm.with.ces$pelas),eta.ces.oly=FPIresults$par,se.eta.ces.oly=se.eta.ces.oly,se.eta.ces.oly.clust=se.eta.ces.oly.clust,eta.ces=eqbm.pre.mm.with.ces$eta[1],n.loops.oly=FPIresults$fpevals))
}

#Doing multiple monte carlo reps for multi-market
monte.oly.mm.rep <- function(parallel=TRUE,settings,v=1,eta.ces.oly.init=3){ #get equilibrium market shares and prices for each market and append them
  if(parallel==T)  foreach(i=1:n.reps, .combine = 'rbind')  %dorng%  
    monte.oly.mm.func(i,settings,v,eta.ces.oly.init) 
  else foreach(i=1:n.reps, .combine = 'rbind')  %do%  {cat(paste("Repetition ",i,"\r"))
    monte.oly.mm.func(i,settings,v,eta.ces.oly.init) }
}


#############################
#   CES OLY approximation   #
#############################
#
#multi-market 
estapprox.oly.mm <- function(pre.input,approx.results,eta.choice) { 
  # CES market share #
  pre <-  as.data.table(pre.input)
  #
  share.ces.oly.mm <- function(results) {
    term <-  pre[, term := exp(results$fitted.values - eta.ces.oly*log(markup) )]$term
    sumterm_n <-  pre[, sumterm_n := sum(term), by=(n)]$sumterm_n
    share_in_n <-  pre[, share_in_n := sum(s), by=(n)]$share_in_n
    pre[,c("term","sumterm_n","share_in_n"):=NULL]
    return((term / sumterm_n)*share_in_n) 
  }
  #
  s.f <- pre$s.f #assumes pre contains s.f
  eta.ces.oly <- eta.choice
  epsilon <- eta.ces.oly + 1 - eta.ces.oly*s.f
  markup <- epsilon/(epsilon-1)
  #
  p.ces.oly <- markup * pre$cit
  s.ces.oly  <- share.ces.oly.mm(results=approx.results)
  return(data.frame(pre,s.ces.oly,p.ces.oly,eta.ces.oly))
}


##################################
# Counterfactual policy changes  #
##################################

#CHANGE to add a meth option.

###### EHA CES OLY #######
CF.approx.oly <- function(pre,post,settings,verbose=T){
  #single destination CF oly EHA, eta.oly is obtained from environment.
  #levels of epsilon
  epsilon <- function(share.firm){
    return(eta.oly+1 - eta.oly*share.firm)
  }
  #exact hat algebra function for Phi oly 
  P.hat <- function(s.init,tau.hat,mu.hat,S0) {
    num <- (tau.hat*mu.hat)^(-eta.oly)
    return((S0+sum(s.init*num))^(-1/eta.oly))
  }
  #exact hat algebra function for CES oly market share
  share.hat <- function(s.init,tau.hat,mu.hat,S0) {
    num <- (tau.hat*mu.hat)^(-eta.oly)
    return(num/(S0+sum(s.init*num)))
  }
  #exact hat algebra function for price elasticity at the f level
  epsilon.hat <- function(s.hat.f){
    return((eta.oly+1 - eta.oly*s.hat.f*df$s.f)/df$eps.f)
  }
  #exact hat algebra of markup
  mu.EHA <- function(eps.hat) {
    num <- eps.hat*df$eps.f
    denom <- df$mu.f*(num-1)
    return(num/denom)    
  }
  #update of share.hat function.
  share.hat.update <- function(s.hat,nu=0.3) {
    dt <- data.table(cbind(df,s.hat))
    dt[,s.f.hat := sum(s.hat*s)/s.f, by=c("f","n")]
    dt[,eps.f.hat := epsilon.hat(s.f.hat)]
    dt[,mu.f.hat := mu.EHA(eps.f.hat)]
    dt[,s.ces.oly.hat := share.hat(s.init=s,tau.hat=tau.hat,mu.hat=mu.f.hat,S0=S.0)]
    return(dt$s.ces.oly.hat)
  }
  # start of eqbm computation #  
  #first load the eta from the environment
  eta.oly <- eta.oly.choice
  #
  df <- data.table(subset(pre))
  #CHANGE HERE TO SUBSTITUTE s.ces.oly for s if(meth == "DEV")
  df[,tau.hat := post$cit/pre$cit]
  df[,s.f := sum(s), by=c("f","n")]
  df[,S.0 := 1 - sum(s), by=c("n")]
  df[,eps.f := epsilon(s.f)]
  df[,mu.f := eps.f/(eps.f-1)]
  #
  s.hat.init <-1
  FPIresults <- fpiter(s.hat.init,fixptfn=share.hat.update,nu=settings$speed[v],control=list(trace=verbose,maxiter=MAXIT))
  n.loops.CF.oly  <- FPIresults$fpevals
  s.ces.oly.hat <- FPIresults$par
  dt <- data.table(cbind(df,s.ces.oly.hat))
  dt[,s.f.ces.oly.hat := sum(s.ces.oly.hat*s)/s.f, by=c("f")]
  dt[,eps.f.hat := epsilon.hat(s.f.ces.oly.hat)]
  dt[,mu.f.hat := mu.EHA(eps.f.hat)]
  dt[,p.ces.oly.hat := mu.f.hat * tau.hat]
  dt[,s.ces.oly := s * s.ces.oly.hat]
  dt[,p.ces.oly := p * p.ces.oly.hat]
  dt[,P.hat := P.hat(s.init=s,tau.hat=tau.hat,mu.hat=mu.f.hat,S0=S.0)] #change in price index
  dt[,profit.ces.oly := (p.ces.oly - post$cit)*((s.ces.oly*Y)/p.ces.oly)]  #model-level profits post-lib (note the post$cit)
  dt[,tar.change := (tau.hat-1) * (pre$cit/pre$mc)] 
  dt[,tarrev.ces.oly := (tar+tar.change) * pre$mc *((s.ces.oly*Y)/p.ces.oly)]
#
  data.frame(post,s.ces.oly=dt$s.ces.oly,p.ces.oly=dt$p.ces.oly,s.ces.oly.hat=dt$s.ces.oly.hat,s.f.ces.oly.hat=dt$s.f.ces.oly.hat,mu.f.hat=dt$mu.f.hat,P.hat=dt$P.hat,profit.ces.oly=dt$profit.ces.oly,tarrev.ces.oly=dt$tarrev.ces.oly,n.loops.CF.oly=n.loops.CF.oly)
}


#for CES oly
CV.ces.oly <- function(pre,post){
  CV.ces <- pre$Y*(1-post$P.hat) # total change in CS
  return(CV.ces[1])
}



#comparison 
CF.compare.oly <- function(pre,post){
  ## market share, considering insiders only:
  pre$s <- pre$s/sum(pre$s)  
  pre$s.ces.oly <- pre$s.ces.oly/sum(pre$s.ces.oly)  
  post$s <- post$s/sum(post$s)  
  post$s.ces.oly <- post$s.ces.oly/sum(post$s.ces.oly)  
  ##
  #
  cr5 <- CR5.agg(pre)
  cor.shr.pre <- cor(pre$s,pre$s.ces.oly)
  reg.results <- lm(pre$s.ces.oly~pre$s)
  reg.shr.lev <- as.numeric(reg.results$coefficients[2])
  #
  s.delta <- post$s - pre$s   
  s.ces.oly.delta <- post$s.ces.oly - pre$s
  # compare sum of changes for domestic firms:
  cor.shr.chg <- cor(s.delta,s.ces.oly.delta) # correlation across models in changes
  reg.results <- lm(s.ces.oly.delta~s.delta)
  reg.shr.dif <- as.numeric(reg.results$coefficients[2])
  dom.shr <- with(pre,sum(s[i==n])) # initial domestic share 
  #mkt share change in pp
  change <- with(post,sum(s[i==n]))-dom.shr 
  change.ces.oly <- with(post,sum(s.ces.oly[i==n]))-dom.shr 
  #PT
  pt.model <- (post$p[post$i!=post$n]-pre$p[pre$i!=pre$n])/(post$cit[post$i!=post$n]- pre$cit[pre$i!=pre$n])
  pte.model <- pt.model * (pre$cit[pre$i!=pre$n]/pre$p[pre$i!=pre$n])
  pt <- mean(pt.model[pt.model!= -Inf & pt.model!= Inf]) # pass thru rate (IO)
  pte <- mean(pte.model[pte.model!= -Inf & pte.model!= Inf]) #  pass thru elasticity  a la Int'l Macro (IM)
  s.f <- pre$s
  s.f[pre$i==pre$n] <- 0 #among foreign models
  s.f[post$cit-pre$cit==0] <-0 # among foreign models that have some variation in tariffs.
  r.f <- rank(-s.f)
  pte1 <- ((post$p[r.f==1]-pre$p[r.f==1])/(post$cit[r.f==1]- pre$cit[r.f==1]))*(pre$cit[r.f==1]/pre$p[r.f==1]) #  pass thru elasticity  a la Int'l Macro (IM) for top foreign firm NOT A MEAN
  #
  #Price elasticity (analytical)
  price.elas <- with(pre,mean(pelas.m))  # mean (across models) of own price elasticity
  #
  converge.prob <- max(pre$n.loops)==MAXIT | max(post$n.loops)==MAXIT | max(post$n.loops.CF.oly) == MAXIT
  #
  return(data.frame(eta=eta.oly.choice,ISE=pre$ISE[1],cr5,pt,pte,pte1,price.elas,cor.shr.pre,reg.shr.lev,cor.shr.chg,reg.shr.dif,dom.shr,change,change.ces=change.ces.oly,converge.prob))
}



#Monte Carlo repetitions for MM
monte.func.cf.oly.mm <- function(repi,settings,v=1,market=1) { # monte carlo stage
  data.exog.mm <- DGP.exog(settings=settings,v=v,grid=subset(tau.grid.allv,vtau==v,select=-vtau)) #gen exog data
  #pre:
  eqbm.pre.mm  <- as.data.table(DGP.endog(parallel=F,v=v,tar.change=0,data.exog = data.exog.mm))
  eqbm.pre.mm.with.ces <- as.data.table(estapprox.mm(eqbm.pre.mm)) #add ces approximation results 
  eqbm.pre.mm.with.ces[, s.f := sum(s), by=c("f","n")]
  eta.ces.oly.init <<- 3
  FPIresults <- fpiter(eta.ces.oly.init,fixptfn=ces.oly.mm.update,pre=eqbm.pre.mm.with.ces,nu=1) #estimate eta.ces.oly
  eta.ces.oly <- FPIresults$par 
  eta.ces <- eqbm.pre.mm.with.ces$eta[1] 
  eta.oly.choice <<- eta.ces.oly
  eqbm.pre.with.ces.oly <- as.data.table(estapprox.oly.mm(eqbm.pre.mm.with.ces,approx.results=ols.results,eta.choice =  eta.oly.choice)) #add ces.oly approximation results
  #post:
  eqbm.post <- eqbm(dat=data.exog.mm,market=market,v=v,tar.change=1) #compute eqbm on one market after removing tariffs
  eqbm.pre.with.ces <- eqbm.pre.mm.with.ces[n==market,] #keep only one destination
  eqbm.pre.with.ces.oly <- eqbm.pre.with.ces.oly[n==market,] #keep only one destination
  eqbm.pre <- eqbm.pre.mm[n==market,] #keep only one destination
  #AFTER THIS, nothing in the code is changed compared to 2ctry, since we have only one destination left.
  eqbm.post.with.ces <- as.data.table(CF.approx(eqbm.pre.with.ces,eqbm.post)) #add ces approximation results
  eqbm.post.with.ces.oly <- as.data.table(CF.approx.oly(pre=eqbm.pre.with.ces,post=eqbm.post.with.ces,settings=settings,verbose = FALSE))
  #
  #compare welfare elements
  CV <- CV(dat=data.exog.mm,market=market,pre=eqbm.pre,post=eqbm.post) #total change in CS
  profit <- sum(eqbm.post$profit[eqbm.post$i==market]-eqbm.pre$profit[eqbm.pre$i==market]) # total change in domestic firms' profits 
  tarrev <- sum(eqbm.post$tarrev[eqbm.post$n!=eqbm.post$i]-eqbm.pre$tarrev[eqbm.pre$n!=eqbm.pre$i]) # total change in tariff revenues
  welfare <- CV+profit+tarrev
  CV.ces <- CV.ces.oly(pre=eqbm.pre.with.ces.oly,post=eqbm.post.with.ces.oly) 
  profit.ces <- sum(eqbm.post.with.ces.oly$profit.ces.oly[eqbm.post.with.ces.oly$i==market]-eqbm.pre$profit[eqbm.pre$i==market]) # total change in domestic firms' profits 
  #tarrev.ces <- tarrev
  tarrev.ces <- sum(eqbm.post.with.ces.oly$tarrev.ces.oly[eqbm.post.with.ces.oly$n!=eqbm.post.with.ces.oly$i]-eqbm.pre$tarrev[eqbm.pre$n!=eqbm.pre$i])
  welfare.ces <- CV.ces+profit.ces+tarrev.ces
  #transform welfare in pct terms
  Y <- eqbm.pre.with.ces$Y[1]
  CV <- CV/Y
  CV.ces <- CV.ces/Y
  profit <- profit / Y
  profit.ces <- profit.ces / Y
  tarrev <- tarrev/Y
  tarrev.ces <- tarrev.ces/Y
  welfare <- welfare / Y
  welfare.ces <- welfare.ces / Y
  #
  return(data.frame(repi=repi,v=v,CF.compare.oly(eqbm.pre.with.ces.oly,eqbm.post.with.ces.oly),CV,profit,tarrev,welfare,CV.ces,profit.ces,tarrev.ces,welfare.ces))
}

#Doing multiple monte carlo reps for multiple markets
monte.rep.cf.oly.mm <- function(parallel=TRUE,settings,v=1,market=1){ #get equilibrium market shares and prices for each market and append them
  if(parallel==T)  foreach(i=1:n.reps, .combine = 'rbind')  %dorng%  monte.func.cf.oly.mm(i,settings,v,market) else foreach(i=1:n.reps, .combine = 'rbind')  %do%  monte.func.cf.oly.mm(i,settings,v,market) 
}

#create model-level data with pre and post market shares for graphs
graph.model.level.func.mm.oly <- function(settings,v=1,market=1) { 
  data.exog.mm <- DGP.exog(settings=settings,v=v,grid=subset(tau.grid.allv,vtau==v,select=-vtau)) #gen exog data
  #pre:
  eqbm.pre.mm  <- as.data.table(DGP.endog(parallel=T,v=v,tar.change=0,data.exog = data.exog.mm))
  eqbm.pre.mm[, s.f := sum(s), by=c("f","n")]
  eta.ces.oly.init <<- 3
  FPIresults <- fpiter(eta.ces.oly.init,fixptfn=ces.oly.mm.update,pre=eqbm.pre.mm,nu=1) #estimate eta.ces.oly
  eta.ces.oly <- FPIresults$par 
  eta.oly.choice <<- eta.ces.oly
  eqbm.pre.oly.mm <- as.data.table(estapprox.oly.mm(eqbm.pre.mm,approx.results=ols.results,eta.choice =  eta.oly.choice)) #add ces.oly approximation results
  #post:
  post <- eqbm(dat=data.exog.mm,market=market,v=v,tar.change= 1) #compute eqbm on one market after removing tariffs
  pre <- eqbm.pre.mm[n==market,] #keep only one destination
  pre.with.ces.oly <- eqbm.pre.oly.mm[n==market,] #keep only one destination
  post.with.ces.oly <- as.data.table(CF.approx.oly(pre=pre,post=post,settings=settings,verbose = FALSE))
  #compare
  s.post <- post.with.ces.oly$s
  s.ces.post <- post.with.ces.oly$s.ces.oly
  s.delta <-s.post-pre$s
  s.ces.delta <-s.ces.post-pre$s
  p.post <- post.with.ces.oly$p
  cit.post <- post.with.ces.oly$cit
  avg.inc.m.post <- post.with.ces.oly$avg.inc.m
  #compute PT
  pte <- ((post$p-pre$p)/(post$cit-pre$cit))*(pre$cit/pre$p) #  pass thru elasticity  a la Int'l Macro (IM)
  rm(pre,post)
  return(data.table(v=v,pre.with.ces.oly,s.post,s.ces.post,s.delta,s.ces.delta,pte,p.post,cit.post,avg.inc.m.post))
}





